import librosa
import numpy as np
import keras
from keras.models import Sequential
from keras.layers import Dense
import math
# from keras_radam import RAdam
# import lookahead

print('读取训练信号\n')
for x in range(500):
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s1\S1_'+str(x+1)+'.wav', sr=None)
    y = y.reshape(1, len(y))
    if(x == 0):
        maleSignal = y
    else:
        maleSignal = np.c_[maleSignal, y]
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s3\S3_' + str(x + 1) + '.wav', sr=None)
    y = y.reshape(1, len(y))
    if (x == 0):
        male2Signal = y
    else:
        male2Signal = np.c_[male2Signal, y]
'#18608750 26857500'
print(len(maleSignal[0]), len(male2Signal[0]))
maleSignal = np.array(maleSignal[0][0:18432000]).reshape(1, -1)
male2Signal = np.array(male2Signal[0][0:18432000]).reshape(1, -1)
print(len(maleSignal[0]))

print("读取测试信号\n")
for x in range(10):
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s1\S1_'+str(x+601)+'.wav', sr=None)
    y = y.reshape(1, len(y))
    if(x == 0):
        cMaleSignal = y
    else:
        cMaleSignal = np.c_[cMaleSignal, y]
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s3\S3_' + str(x + 601) + '.wav',sr=None)
    y = y.reshape(1, len(y))
    if (x == 0):
        cMale2Signal = y
    else:
        cMale2Signal = np.c_[cMale2Signal, y]
'#373510 517760'
print(len(cMaleSignal[0]), len(cMale2Signal[0]))
cMaleSignal = np.array(cMaleSignal[0][0:363520]).reshape(1, -1)
cMale2Signal = np.array(cMale2Signal[0][0:363520]).reshape(1, -1)
print(len(cMaleSignal[0]), len(cMaleSignal[0]) == len(cMale2Signal[0]))
cMixedSignal = np.array(cMaleSignal[0]+cMale2Signal[0]).reshape(1, -1)


print("信号变换：")
#训练男生2信号
bMale2Signal = librosa.stft(male2Signal[0], n_fft=512)
bMale2SignalAmplitude = np.abs(bMale2Signal)
bMale2SignalAngle = np.angle(bMale2Signal)
#训练男生
bMaleSignal = librosa.stft(maleSignal[0], n_fft=512)
bMaleSignalAmplitude = np.abs(bMaleSignal)
bMaleSignalAngle = np.angle(bMaleSignal)
#训练  目标IRM
IRM = bMaleSignalAmplitude/(bMaleSignalAmplitude+bMale2SignalAmplitude)
TrainIRM = IRM.transpose()
#训练输入混合信号
bMixedSignalAmplitude = bMale2SignalAmplitude+bMaleSignalAmplitude
bMixedSignalMax = np.max(bMixedSignalAmplitude, axis=0)
bMixedSignalNorm = bMixedSignalAmplitude/bMixedSignalMax
TrainMixedSignal = bMixedSignalNorm.transpose()
#测试混合信号
tMixedSignal = librosa.stft(cMixedSignal[0], n_fft=512)
tMixedSignalAmplitude = np.abs(tMixedSignal)
tMixedSignalAngle = np.angle(tMixedSignal)
tMixedSignalAmplitudeMax = np.max(tMixedSignalAmplitude, axis=0)
tMixedSignalAmplitudeNorm = tMixedSignalAmplitude/tMixedSignalAmplitudeMax
TestMixedSignal = tMixedSignalAmplitudeNorm.transpose()
print("信号变换完成")

print("训练开始：")
X_Train = TrainMixedSignal
Y_Train = TrainIRM
X_Test = TestMixedSignal
model = Sequential()
model.add(Dense(1024, activation='relu', input_dim=257))
model.add(Dense(1024, activation='selu'))
model.add(Dense(1024, activation='sigmoid'))
model.add(Dense(257, activation='relu'))
sgd = keras.optimizers.SGD(lr=0.01, momentum=0.9)
NADAM =keras.optimizers.Nadam(lr=0.002, beta_1=0.9, beta_2=0.999)
model.compile(optimizer="adam", loss='mse', metrics=['accuracy'])
# lookahead1 = lookahead.Lookahead(k=5, alpha=0.5)
# lookahead1.inject(model)
model.fit(X_Train, Y_Train, epochs=200, batch_size=256, verbose=2)
Y_pred = model.predict(X_Test, batch_size=256, verbose=0, steps=None)
print("训练完成;")

print("语音输出：")
rMaleSignalAmplitude = (Y_pred*TestMixedSignal).transpose()*tMixedSignalAmplitudeMax
rMaleSignal = rMaleSignalAmplitude*np.power(math.e, 1j*tMixedSignalAngle)
rebuildMaleSignal = librosa.istft(rMaleSignal)
librosa.output.write_wav(r".\speech\m2\IRM\adam\rebuild_signal.wav", rebuildMaleSignal, sr=25000)
librosa.output.write_wav(r".\speech\m2\IRM\adam\origin_signal.wav", cMaleSignal[0], sr=25000)
print("语音输出完成")
